from asyncio import FastChildWatcher
import time
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import torch
import cv2

from draw import *

from pcgworker.PCGWorker import *

from wfc_env import *

'''
simulatiom parameters
'''

num_matrix_envs = 1
headless = False
spd = 5
max_actors_num_per_env = 9*9*4+1

# get default set of parameters
sim_params = gymapi.SimParams()
args = gymutil.parse_arguments()
sim_params.physx.solver_type = 1
sim_params.physx.num_position_iterations = 2
sim_params.physx.num_velocity_iterations = 1
sim_params.physx.num_threads = args.num_threads
sim_params.physx.use_gpu = args.use_gpu
sim_params.use_gpu_pipeline = False #args.use_gpu_pipeline
sim_params.substeps = 2
sim_params.dt = 1.0 / 30.0
# setting up the Z-up axis
sim_params.up_axis = gymapi.UP_AXIS_Z
sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.8)

# configure the ground plane
plane_params = gymapi.PlaneParams()
plane_params.normal = gymapi.Vec3(0, 0, 1) # z-up!
plane_params.distance = 1
plane_params.static_friction = 1
plane_params.dynamic_friction = 1
plane_params.restitution = 0.1

# set up the env grid
num_envs = 1
spacing = 20
env_lower = gymapi.Vec3(-spacing, -spacing, -spacing)
env_upper = gymapi.Vec3(spacing, spacing, spacing)

# Attractor setup
attractor_properties = gymapi.AttractorProperties()
attractor_properties.stiffness = 5e10
attractor_properties.damping = 0
attractor_properties.axes = gymapi.AXIS_ROTATION
attractor_pose = gymapi.Transform()
attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
attractor_pose.r = gymapi.Quat(0,1,0,1)

# create procedural asset
asset_options = gymapi.AssetOptions()
asset_options.density = 1.0
asset_options.fix_base_link = False
asset_options1 = gymapi.AssetOptions()
asset_options1.density = 0.02
asset_options1.fix_base_link = True
asset_options1.linear_damping = 0.

color_red = gymapi.Vec3(1,0,0)

# load PCG asset
asset_root = "./assets"
pt_gray_cube = "PCG/gray_cube.urdf"
pt_blue_cube = "PCG/blue_cube.urdf"
pt_red_cube = "PCG/red_cube.urdf"
pt_white_cube = "PCG/white_cube.urdf"
pt_yellow_cube = "PCG/yellow_cube.urdf"
pt_blue_ramp = "PCG/blue_ramp.urdf"
pt_red_ramp = "PCG/red_ramp.urdf"
pt_white_ramp = "PCG/white_ramp.urdf"
pt_yellow_ramp = "PCG/yellow_ramp.urdf"

# ramp and corner attitutde
ramp_rot0 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0*math.pi)
ramp_rot1 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi)
ramp_rot2 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 1*math.pi)
ramp_rot3 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 1.5*math.pi)
corner_rot0 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0*math.pi)
corner_rot1 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi)
corner_rot2 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 1*math.pi)
corner_rot3 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 1.5*math.pi)

# birick size
side_l = 2

# camera sensor properties
camera_properties = gymapi.CameraProperties()
camera_properties.width = WIDTH
camera_properties.height = HEIGHT
camera_properties.enable_tensors = True

# rigid shape material properties
shape_props = [gymapi.RigidShapeProperties()]
shape_props[0].friction = -1
shape_props[0].rolling_friction = 1
shape_props[0].torsion_friction = 1
shape_props[0].compliance = 0
shape_props[0].restitution = 0
shape_props1 = [gymapi.RigidShapeProperties()]
shape_props1[0].friction = 0.2
shape_props1[0].rolling_friction = -1
shape_props1[0].torsion_friction = -1
shape_props1[0].compliance = 0
shape_props1[0].restitution = 0
shape_props2 = [gymapi.RigidShapeProperties()]
shape_props2[0].friction = 1
shape_props2[0].rolling_friction = -1
shape_props2[0].torsion_friction = -1
shape_props2[0].compliance = 0
shape_props2[0].restitution = 0

'''
generate WFC maps
'''
PCGWorker_ = PCGWorker(9,9)
seed = PCGWorker_.generate()
seeds = [seed]
for i in range(0,num_matrix_envs-1):
    seed_ = PCGWorker_.mutate(seeds[-1],81)
    seeds.append(seed_)
building_block_pointer = [0,0,0,0,0,0,0,0,0,0,0,0,0]    # 0-4,brick; 5-8,corners; 9-12 ,ramps;

gym = gymapi.acquire_gym()

sim = gym.create_sim(args.compute_device_id, args.graphics_device_id, gymapi.SIM_PHYSX, sim_params)

print("args.graphics_device_id : ",args.graphics_device_id)
print("args.compute_device_id : ",args.compute_device_id)


# create the ground plane
gym.add_ground(sim, plane_params)

capsule_asset = gym.create_capsule(sim, 1, 1, asset_options)

asset_gray_cube = gym.load_asset(sim, asset_root, pt_gray_cube,asset_options1)
asset_blue_cube = gym.load_asset(sim, asset_root, pt_blue_cube,asset_options1)
asset_red_cube = gym.load_asset(sim, asset_root, pt_red_cube,asset_options1)
asset_white_cube = gym.load_asset(sim, asset_root, pt_white_cube,asset_options1)
asset_yellow_cube = gym.load_asset(sim, asset_root, pt_yellow_cube,asset_options1)
asset_blue_ramp = gym.load_asset(sim, asset_root, pt_blue_ramp,asset_options1)
asset_red_ramp = gym.load_asset(sim, asset_root, pt_red_ramp,asset_options1)
asset_white_ramp = gym.load_asset(sim, asset_root, pt_white_ramp,asset_options1)
asset_yellow_ramp = gym.load_asset(sim, asset_root, pt_yellow_ramp,asset_options1)

box_asset = gym.create_box(sim, 1, 1, 1, asset_options)

if not headless:
    cam_props = gymapi.CameraProperties()
    viewer = gym.create_viewer(sim, cam_props)

envs = []
camera_handles = []
actor_handles = []
attractor_handles = []

for n in range(num_matrix_envs):
    # create environments
    env_ = gym.create_env(sim, env_lower, env_upper, 4)
    envs.append(env_)

    n_actors = 0
    # create actors based on WFC seed
    for i in range(0,9):
        for j in range(0,9):
            tile_ = seeds[n].wave_oriented[j*9+i][0][0]
            rot = seeds[n].wave_oriented[j*9+i][0][1]
            pose_ij = gymapi.Transform()
            pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 0*side_l)
            pose_ij.r = gymapi.Quat(0,0,0,1)
            actor = gym.create_actor(env_, asset_gray_cube, pose_ij, "gray_cube" + str(i) + str(j), n, 0)
            # cubes
            if tile_<=5:
                if tile_==1:
                    pass
                elif tile_==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                elif tile_==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                elif tile_==4:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    actor = gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                elif tile_==5:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    actor = gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    actor = gym.create_actor(env_, asset_white_cube, pose_ij, "white_cube" + str(i) + str(j), n, 0)
            # corners
            elif tile_==6:
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = corner_rot2
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = corner_rot1
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = corner_rot0
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = corner_rot3
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
            elif tile_==7:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = corner_rot2
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = corner_rot1
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = corner_rot0
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = corner_rot3
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
            elif tile_==8:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = corner_rot2
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = corner_rot1
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = corner_rot0
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = corner_rot3
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
            elif tile_==9:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                actor = gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = corner_rot2
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = corner_rot1
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = corner_rot0
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = corner_rot3
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
            # ramps
            elif tile_==10:
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = ramp_rot0
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = ramp_rot1
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = ramp_rot2
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                    pose_ij.r = ramp_rot3
                    actor = gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
            elif tile_==11:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = ramp_rot0
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = ramp_rot1
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = ramp_rot2
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                    pose_ij.r = ramp_rot3
                    actor = gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
            elif tile_==12:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = ramp_rot0
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = ramp_rot1
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = ramp_rot2
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                    pose_ij.r = ramp_rot3
                    actor = gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
            elif tile_==13:
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 1*side_l)
                actor = gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 2*side_l)
                actor = gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 3*side_l)
                actor = gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                if rot==0:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = ramp_rot0
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==1:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = ramp_rot1
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==2:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = ramp_rot2
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                elif rot==3:
                    pose_ij.p = gymapi.Vec3(i*side_l, j*side_l, 4*side_l)
                    pose_ij.r = ramp_rot3
                    actor = gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                gym.set_actor_rigid_shape_properties(env_, actor, shape_props)

    # create capsule_asset actor in the environment
    pose = gymapi.Transform()
    pose.p = gymapi.Vec3(15, 15, 30)
    pose.r = gymapi.Quat(0,1,0,1)
    cap_handle = gym.create_actor(env_, capsule_asset, pose, 'capsule_asset', n, 0)
    gym.set_actor_rigid_shape_properties(env_, cap_handle, shape_props2)

    gym.set_actor_scale(env_, cap_handle, 0.4)
    actor_handles.append(cap_handle)

    attractor_properties_ = gymapi.AttractorProperties()
    attractor_properties_.stiffness = 5e10
    attractor_properties_.damping = 0
    attractor_properties_.axes = gymapi.AXIS_ROTATION
    attractor_properties_.target = pose
    attractor_properties_.rigid_handle = cap_handle
    attractor_handle_ = gym.create_rigid_body_attractor(env_, attractor_properties_)
    gym.set_rigid_body_color(env_, cap_handle, 0, gymapi.MESH_VISUAL_AND_COLLISION, color_red)
    
    h1 = gym.create_camera_sensor(env_, camera_properties)
    camera_offset = gymapi.Vec3(-2, 0, 0)
    camera_rotation = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 0, 1), 0.4*math.pi)
    body_handle = gym.get_actor_rigid_body_handle(env_, cap_handle, 0)
    gym.attach_camera_to_body(h1, env_, body_handle, gymapi.Transform(camera_offset, camera_rotation), gymapi.FOLLOW_TRANSFORM)
    camera_handles.append(h1)
    attractor_handles.append(attractor_handle_)

    actor_count = gym.get_actor_count(env_)
    # create blank actors
    num_blk_act = max_actors_num_per_env - actor_count
    for i in range(num_blk_act):
        pose = gymapi.Transform()
        pose.p = gymapi.Vec3(0, 0, 0)
        pose.r = gymapi.Quat(0,1,0,1)
        actor = gym.create_actor(env_, box_asset, pose, 'box_asset' + str(i), n, 0)
        gym.set_actor_scale(env_, actor, 0.000001)


# a must-have operation 
gym.prepare_sim(sim)

frame_count = 0

if not headless:
    cam_pos = gymapi.Vec3(0.0, 15.0, 15.0)
    cam_target = gymapi.Vec3(1.5, 4, 0)
    gym.viewer_camera_look_at(viewer, None, cam_pos, cam_target)

# time test
start_time = time.time()
last_frame_cnt = 0

# Create helper geometry used for visualization
# Create an wireframe axis
axes_geom = gymutil.AxesGeometry(2)

facing = 0
facing_step = 0.25

# local coordinate system
verts = np.empty((3, 1), gymapi.Vec3.dtype)
verts[0][0] = (1, 0, 0)
verts[1][0] = (0, 1, 0)
verts[2][0] = (0, 0, 1)

# store initial status for all envs
sim_status = []
for i in range(num_matrix_envs):
    env_status = np.copy(gym.get_env_rigid_body_states(envs[i], gymapi.STATE_ALL))
    sim_status.append(env_status)

while True:

    if not headless:
        if gym.query_viewer_has_closed(viewer):
            break
        # clear lines from previous frame
        gym.clear_lines(viewer)

    '''
    step begin
    '''
    # step the physics
    gym.simulate(sim)
    gym.fetch_results(sim, True)

    # if not headless:
    gym.step_graphics(sim)

    # render the camera sensors
    gym.render_all_camera_sensors(sim)

    for i in range(num_matrix_envs):
        rgb_image = gym.get_camera_image(sim, envs[i], camera_handles[i], gymapi.IMAGE_COLOR)

        if not headless :
            # create an opencv image of size (HEIGHT,WIDTH , 4)
            img = np.zeros((HEIGHT,WIDTH, 4), np.uint8)
            img[:, :, :4] = rgb_image.reshape((HEIGHT,WIDTH, 4))
            # rotate 90 degrees counter-clockwise
            img = np.rot90(img, -1)
            # transform the image from RGBA to RGB
            img = img[:, :, :3]
            # transform the image from RGB to BGR
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

            '''
            render
            '''
            cv2.imshow('rgb', img)
            action = cv2.waitKey(1)

            '''
            step : action
            '''
            if action >0 :
                print(action)
                body_states = gym.get_actor_rigid_body_states(envs[i], actor_handles[i], gymapi.STATE_ALL)
                # transform global coordinate system to local coordinate system
                body_ = gym.get_actor_rigid_body_handle(envs[i], actor_handles[i], 0)
                body_t = gym.get_rigid_transform(envs[i], body_)
                body_t.p = gymapi.Vec3(0,0,0)
                verts_ = body_t.transform_points(verts)
                # print(verts_)

                if action == 32:
                    body_states['vel']['linear'].fill((0,0,spd))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 119:
                    body_states['vel']['linear'].fill((spd * verts_[1][0][0],spd * verts_[1][0][1],spd * verts_[1][0][2]))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 115:
                    body_states['vel']['linear'].fill((-spd * verts_[1][0][0],-spd * verts_[1][0][1],-spd * verts_[1][0][2]))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 97:
                    body_states['vel']['linear'].fill((-spd * verts_[2][0][0],-spd * verts_[2][0][1],-spd * verts_[2][0][2]))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 100:
                    body_states['vel']['linear'].fill((spd * verts_[2][0][0],spd * verts_[2][0][1],spd * verts_[2][0][2]))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)
                elif action == 107:
                    facing -= facing_step
                    attractor_pose = gymapi.Transform()
                    attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
                    attractor_pose.r = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), facing*math.pi)
                    gym.set_attractor_target(envs[i], attractor_handles[i], attractor_pose)
                elif action == 108:
                    facing += facing_step
                    attractor_pose = gymapi.Transform()
                    attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
                    attractor_pose.r = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), facing*math.pi)
                    gym.set_attractor_target(envs[i], attractor_handles[i], attractor_pose)
                # break !!
                elif action == 111:
                    body_states['vel']['linear'].fill((0,0,0))
                    gym.set_actor_rigid_body_states(envs[i], actor_handles[i], body_states, gymapi.STATE_VEL)

                # reset !!
                elif action == 114:
                    print("reset")
                    gym.set_env_rigid_body_states(envs[i], sim_status[i], gymapi.STATE_ALL)

    '''
    step end
    '''

    '''
    render
    '''
    # draw camera entities
    if not headless:
        for i in range(num_matrix_envs):
            body_states = gym.get_camera_transform(sim, envs[i], camera_handles[i])
            draw_camera(gym, viewer,envs[i],body_states)
    if not headless:
        # update the viewer
        gym.draw_viewer(viewer, sim, True)

    frame_count += 1
    end_time = time.time()
    if (end_time - start_time) > 1:
        print("FPS: %.2f" % ((frame_count - last_frame_cnt)*num_matrix_envs))
        last_frame_cnt = frame_count
        start_time = time.time()

if not headless:
    gym.destroy_viewer(viewer)
    
gym.destroy_sim(sim)
